import os

from langchain.chains.combine_documents.map_reduce import MapReduceDocumentsChain
from langchain.chains.combine_documents.reduce import ReduceDocumentsChain
from langchain.chains.combine_documents.stuff import StuffDocumentsChain
from langchain.chains.llm import LLMChain
from langchain_community.document_loaders import WebBaseLoader
from langchain_core.prompts import PromptTemplate
from langchain_openai import ChatOpenAI
from langchain_text_splitters import CharacterTextSplitter
from pydantic import SecretStr

# 读取API密钥
api_key = os.getenv("DASHSCOPE_API_KEY")
if not api_key:
    raise ValueError("请设置环境变量DASHSCOPE_API_KEY（阿里云百炼API-KEY）")

# 创建大语言模型实例
model = ChatOpenAI(
    model="qwen-plus-latest",
    temperature=0.5,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    api_key=SecretStr(api_key),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# 加载文档,使用langchain提供的WebBaseModel进行加载网上博文
loader = WebBaseLoader("https://lilianweng.github.io/posts/2023-06-23-agent/")
docs = loader.load()  # 得到篇文章

# Map-Reduce
# 第一步：切割阶段
text_splitter = CharacterTextSplitter.from_tiktoken_encoder(chunk_size=1024, chunk_overlap=0)
split_docs = text_splitter.split_documents(docs)

# 第二步：map阶段
map_template = """
    以下是一组文档(documents)
    "{docs}"
    根据这个文档列表，简洁的总结摘要:
"""
map_prompt = PromptTemplate.from_template(map_template)
map_llm_chain = LLMChain(llm=model, prompt=map_prompt)

# 第三步：reduce阶段
reduce_template = """
    以下是一组总结摘要
    {docs}
    将这些内容提炼出一个最终的、统一的总结摘要:
"""
reduce_prompt = PromptTemplate.from_template(reduce_template)
reduce_llm_chain = LLMChain(llm=model, prompt=reduce_prompt)
'''
reduce的思路
如果map之后的文档的累计token数超过了 4000个，那么我们将递归的将文档以 <=4000个token 的批次传递给我们的StuffDocumentsChain
一旦这些批量摘要的累计大小 <4000个token，我们将它们全部传递给StuffDocumentsChain最后一次，用来创建最终摘要。
'''

# 第四步：combine阶段
combine_chain = StuffDocumentsChain(llm_chain=reduce_llm_chain, document_variable_name='docs')
reduce_chain = ReduceDocumentsChain(
    combine_documents_chain=combine_chain,  # 最终调用的链
    collapse_documents_chain=combine_chain,  # 中间汇总的链
    token_max=4096  # 将文本分组的最大token数
)

# 第五步：合并所有链
map_reduce_chain = MapReduceDocumentsChain(
    llm_chain=map_llm_chain,
    reduce_documents_chain=reduce_chain,
    document_variable_name='docs',
    return_intermediate_steps=False
)

# 调用
result = map_reduce_chain.invoke(split_docs)
print(result['output_text'])
